import time
import struct
import enum
import logging
import math
from copy import deepcopy
import numpy as np
from serial import Serial
from serial.tools import list_ports
import asyncio
import websockets
import json
import threading

# region: GBot/utils.py

def bytes_to_short(data: bytes, signed: bool = False, byteorder: str = 'little') -> int:
    if len(data) != 2:
        raise ValueError("Data must be exactly 2 bytes long")
    prefix = '<' if byteorder == 'little' else '>'
    format_char = 'h' if signed else 'H'
    return struct.unpack(f"{prefix}{format_char}", data)[0]

def short_to_bytes(value: int, signed: bool = False, byteorder: str = 'little') -> bytes:
    if signed:
        format_char = 'h'
        min_val, max_val = -32768, 32767
    else:
        format_char = 'H'
        min_val, max_val = 0, 65535
    if not (min_val <= value <= max_val):
        raise OverflowError(f"Value {value} out of range for {'signed' if signed else 'unsigned'} short")
    prefix = '<' if byteorder == 'little' else '>'
    return struct.pack(f"{prefix}{format_char}", value)

def bytes_to_int(byte_data, signed=False, byteorder='little'):
    if len(byte_data) != 4:
        raise ValueError("输入必须是 4 字节")
    fmt_char = 'i' if signed else 'I'
    fmt_str = ('>' if byteorder == 'big' else '<') + fmt_char
    return struct.unpack(fmt_str, byte_data)[0]

def int_to_bytes(int_value, signed=False, byteorder='little'):
    if signed and not (-2**31 <= int_value < 2**31):
        raise ValueError("有符号整数超出 4 字节范围")
    elif not signed and not (0 <= int_value < 2**32):
        raise ValueError("无符号整数超出 4 字节范围")
    fmt_char = 'i' if signed else 'I'
    fmt_str = ('>' if byteorder == 'big' else '<') + fmt_char
    return struct.pack(fmt_str, int_value)

# endregion

# region: GBot/global_state.py

class Address(enum.Enum):
    DEVICE_UUID         = (0, 4)
    VERSION             = (4, 2)
    MOTOR_TYPE          = (6, 1)
    CURRENT_POSITION    = (7, 2)
    CURRENT_SPEED       = (9, 2)
    CURRENT_LOAD        = (11, 2)
    CURRENT_VOLTAGE     = (13, 1)
    CURRENT_CURRENT     = (14, 2)
    CURRENT_TEMPERATURE = (16, 1)
    TORQUE_ENABLE       = (50, 1)
    TARGET_POSITION     = (51, 2)
    ID                  = (70, 1)
    MIN_POSITION        = (71, 2)
    MAX_POSITION        = (73, 2)
    POSITION_OFFSET     = (75, 2)
    MAX_VOLTAGE         = (77, 1)
    MIN_VOLTAGE         = (78, 1)
    MAX_TEMPERATURE     = (79, 1)
    MAX_CURRENT         = (80, 2)
    KP                  = (82, 1)
    KI                  = (83, 1)
    KD                  = (84, 1)

    @classmethod
    def get_address(cls, address:int):
        for addr in cls:
            if addr.value[0] == address:
                return addr
        return None

class ErrorCode(enum.Enum):
    SUCCESS             = 0
    WRITE_ERROR         = 1
    READ_ERROR          = 2
    READ_TIMEOUT        = 3

class Result:
    def __init__(self, error: ErrorCode = ErrorCode.SUCCESS, frame: list[int] = None, input = None):
        self.__error_code = error
        self.__frame = frame
        self.__input = input
        self.__value_map = {}

        if frame is None or input is None:
            return

        id = frame[2]
        cmd = frame[3]
        if cmd != 0x03:
            return
        if id != 0xFF and id < 128 and id >= 248:
            return

        addresses = []
        if isinstance(input, Address):
            addresses.append(input)
        elif isinstance(input, list):
            addresses.extend(input)

        cnt = 6 if id == 0xFF else 5

        while cnt < len(frame) - 2:
            addr = Address.get_address(frame[cnt])
            if addr is None:
                break
            addr_int = addr.value[0]
            addr_len = addr.value[1]

            if addr_len == 1:
                self.__value_map[addr_int] = frame[cnt+1]
            elif addr_len == 2:
                self.__value_map[addr_int] = bytes_to_short(bytearray(frame[cnt+1:cnt+3]))
            elif addr_len == 4:
                self.__value_map[addr_int] = bytes_to_int(bytearray(frame[cnt+1:cnt+5]))
            cnt += addr_len + 1

    def is_success(self) -> bool:
        return self.__error_code == ErrorCode.SUCCESS

    def get_error_code(self) -> int:
        return self.__error_code.value

    def get_data(self, address: Address) -> int:
        address_int = address.value[0]
        return self.__value_map.get(address_int)

# endregion

# region: GBot/port_handler.py

class PortHandler:
    def __init__(self):
        self.__serial: Serial = None
        self._port = None
        self._baudrate = 230400
        self._bytesize = 8
        self._parity = 'N'
        self._stopbits = 1
        self._read_timeout = None
        self._write_timeout = None
        self.__is_running = False

    @property
    def baudrate(self):
        return self._baudrate

    @baudrate.setter
    def baudrate(self, value):
        if self.__serial and self.__serial.is_open:
            raise ValueError("无法修改已打开的串口波特率")
        self._baudrate = value

    def open(self, port) -> bool:
        self.close()
        try:
            self._port = port
            self.__serial = Serial(port=port, baudrate=self._baudrate, bytesize=self._bytesize, parity=self._parity, stopbits=self._stopbits, timeout=self._read_timeout, write_timeout=self._write_timeout)
            self.__is_running = True
            return True
        except Exception:
            return False

    def is_open(self) -> bool:
        return self.__serial and self.__serial.is_open

    def close(self):
        if self.__serial and self.__serial.is_open:
            self.__serial.close()
            self.__is_running = False
            self.__serial = None

    def read_port(self, length:int):
        if self.__serial and self.__serial.is_open:
            return self.__serial.read(length)

    def write_port(self, data):
        if self.__serial and self.__serial.is_open:
            self.__serial.reset_input_buffer()
            self.__serial.write(data)
            self.__serial.flush()

    def in_waiting(self):
        if self.__serial and self.__serial.is_open:
            return self.__serial.in_waiting
        return 0

# endregion

# region: GBot/sync_connector.py

FRAME_HEADER    = 0xAA
FRAME_TAIL      = 0xBB
FRAME_CMD_READ          = 0x03

def checksum(id: int, cmd: int, data: list[int]) -> int:
    return (id + cmd + len(data) + sum(data)) & 0xFF

def frame_generator(id: int, cmd: int, data: list[int]) -> bytearray:
    frame = bytearray()
    frame.append(FRAME_HEADER)
    frame.append(FRAME_HEADER)
    frame.append(id)
    frame.append(cmd)
    frame.append(len(data))
    for d in data:
        frame.append(d)
    frame.append(checksum(id, cmd, data))
    frame.append(FRAME_TAIL)
    return frame

class SyncConnector:
    def __init__(self, portHandler: PortHandler):
        self.__port_handler = portHandler

    def _process_address_data(self, address, data_value):
        data = []
        if isinstance(address, list):
            if not isinstance(data_value, list) or len(address) != len(data_value):
                raise ValueError("地址和值的数量必须匹配")
            for i in range(len(address)):
                reg = address[i].value[0]
                value_len = address[i].value[1]
                data.extend([reg, value_len])
                data.extend(self._value_to_bytes(data_value[i], value_len))
        else:
            reg = address.value[0]
            value_len = address.value[1]
            data.extend([reg, value_len])
            data.extend(self._value_to_bytes(data_value, value_len))
        return data

    def _value_to_bytes(self, value, length):
        if length == 1:
            return [value & 0xFF]
        elif length == 2:
            return short_to_bytes(value)
        elif length == 4:
            return int_to_bytes(value)
        raise ValueError("无效的值长度")

    def _parse_response_frame(self) -> Result:
        retry_cnt = 0
        read_list = []
        state = 0
        self.__port_handler._read_timeout = 1
        while True:
            in_waiting = self.__port_handler.in_waiting()
            if in_waiting == 0:
                if retry_cnt < 5:
                    retry_cnt += 1
                    time.sleep(0.01)
                    continue
                else:
                    state = -1
                    break
            read_list.extend(list(self.__port_handler.read_port(in_waiting)))
            while len(read_list) >= 7:
                if read_list[0] != FRAME_HEADER or read_list[1] != FRAME_HEADER:
                    read_list.pop(0)
                    continue
                data_length = read_list[4]
                if data_length > 48 or len(read_list) < 7 + data_length or read_list[6 + data_length] != FRAME_TAIL:
                    read_list.pop(0)
                    continue
                checksum_val = sum(read_list[2:5 + data_length]) & 0xFF
                if checksum_val != read_list[5 + data_length]:
                    read_list.pop(0)
                    continue
                read_list = read_list[0:7 + data_length]
                state = 1
                break
            if state == 1:
                break
        if state == -1:
            return Result(error=ErrorCode.READ_TIMEOUT)
        return Result(frame=read_list, input=self.last_read_address)

    def read(self, id_list: list[int], address_list: list[Address]) -> Result:
        self.last_read_address = address_list
        data = []
        for address in address_list:
            data.extend([address.value[0], address.value[1]])
        frame = frame_generator(id_list[0], FRAME_CMD_READ, data)
        self.__port_handler.write_port(frame)
        return self._parse_response_frame()

    def write(self, id_list, address, data_value):
        data = self._process_address_data(address, data_value)
        if isinstance(id_list, list):
            # Simplified broadcast write for multiple motors
            frame = frame_generator(id_list[0], 0x01, data) 
        else:
            frame = frame_generator(id_list, 0x01, data)
        self.__port_handler.write_port(frame)
        # Since this is a write operation, we don't expect a response frame in this simplified model
        return Result(error=ErrorCode.SUCCESS)

# endregion

# region: gbot.py

class RobotDeviceAlreadyConnectedError(Exception):
    pass

class RobotDeviceNotConnectedError(Exception):
    pass

class GBotMotorsBus:
    def __init__(self, port: str, motors: dict[str, tuple[int, str]], mock=False):
        self.port = port
        self.motors = motors
        self.mock = mock
        self.__port_handler: PortHandler = PortHandler()
        self.__sync_connector: SyncConnector = SyncConnector(self.__port_handler)
        self.is_connected = False
        self.calibration = {}
        self.address_mapping = {
            'PRESENT_POSITION': 'CURRENT_POSITION',
        }

    def _get_address_from_item(self, item: str):
        item_upper = item.upper()
        mapped_item = self.address_mapping.get(item_upper, item_upper)
        try:
            return Address[mapped_item]
        except KeyError:
            raise KeyError(f"Control table item '{item}' not found. It was resolved to '{mapped_item}' which is not a valid Address member.")

    def connect(self):
        if self.is_connected:
            raise RobotDeviceAlreadyConnectedError(f"GBotMotorsBus is already connected to port {self.port}.")
        if not self.__port_handler.open(self.port):
            raise RobotDeviceNotConnectedError(f"Failed to connect to port {self.port}.")
        self.is_connected = True

    def disconnect(self):
        if not self.is_connected:
            raise RobotDeviceNotConnectedError(f"GBotMotorsBus is not connected to port {self.port}.")
        self.__port_handler.close()
        self.is_connected = False

    @property
    def motor_names(self) -> list[str]:
        return list(self.motors.keys())

    @property
    def motor_indices(self) -> list[int]:
        return [idx for idx, _ in self.motors.values()]

    def set_calibration(self, calibration: dict[str, list]):
        self.calibration = calibration

    def read(self, item: str):
        address = self._get_address_from_item(item)
        motor_ids = self.motor_indices
        positions = []
        for motor_id in motor_ids:
            result = self.__sync_connector.read([motor_id], [address])
            if result.is_success():
                pos = result.get_data(address)
                if pos is not None:
                    positions.append(pos)
                else:
                    positions.append(0)
            else:
                positions.append(0)
        return np.array(positions)

    def write(self, item: str, value):
        address = self._get_address_from_item(item)
        motor_ids = self.motor_indices
        # Simplified broadcast write
        self.__sync_connector.write(motor_ids, address, value)

# endregion

# region: read_teacher.py main logic

# --- WebSocket Server Setup ---
clients = set()
latest_joint_data = {}

async def handler(websocket):
    clients.add(websocket)
    try:
        await websocket.wait_closed()
    finally:
        clients.remove(websocket)

async def broadcast_joint_data():
    while True:
        if clients and latest_joint_data:
            message = json.dumps(latest_joint_data)
            print(f"Broadcasting: {message}") # Debug print
            # Use asyncio.gather to run sends concurrently, ignoring exceptions for disconnected clients
            await asyncio.gather(*[client.send(message) for client in clients], return_exceptions=True)
        await asyncio.sleep(0.1) # Broadcast every 100ms

async def main_server_logic():
    # Start the broadcast task that runs in the background
    broadcast_task = asyncio.create_task(broadcast_joint_data())
    
    # Start the WebSocket server and run it forever
    async with websockets.serve(handler, "localhost", 8765, ping_interval=20, ping_timeout=20):
        await asyncio.Future()  # This will run until cancelled

def run_websocket_server():
    asyncio.run(main_server_logic())

# --- Main Application Logic ---

def main():
    # Start the WebSocket server in a separate thread
    ws_thread = threading.Thread(target=run_websocket_server, daemon=True)
    ws_thread.start()

    port = 'COM4'
    motors = {
        'joint_1': (1, 'genki_M1'),
        'joint_2': (2, 'genki_M1'),
        'joint_3': (3, 'genki_M1'),
        'joint_4': (4, 'genki_M1'),
        'joint_5': (5, 'genki_M1'),
        'joint_6': (6, 'genki_M1'),
    }

    # Homing offset for zero position, assuming 2048 is the center for a 4096 resolution.
    homing_offset = 2048
    # Motor resolution (steps per revolution)
    resolution = 4096

    bus = GBotMotorsBus(port=port, motors=motors)

    try:
        while True:  # Main reconnection loop
            try:
                print(f"Connecting to port {bus.port}...")
                bus.connect()
                print("Connection successful.")

                print("Enabling torque for all motors...")
                bus.write('Torque_Enable', 1)
                time.sleep(0.5)

                print("Start reading motor positions. Press Ctrl+C to stop.")
                global latest_joint_data
                while True:  # Reading loop
                    positions = bus.read('Present_Position')
                    min_positions = bus.read('MIN_POSITION')
                    max_positions = bus.read('MAX_POSITION')

                    raw_positions_dict = {}
                    for i, name in enumerate(bus.motor_names):
                        pos = positions[i]
                        min_pos = min_positions[i]
                        max_pos = max_positions[i]
                        
                        # print(f"Motor '{name}' (ID: {bus.motor_indices[i]}) | Raw Pos: {pos} | Min Pos: {min_pos} | Max Pos: {max_pos}")

                        # New conversion logic based on user feedback:
                        # 2048 -> 0 degrees
                        # 1024 -> -90 degrees
                        # 3072 -> +90 degrees
                        # This means a range of 1024 units corresponds to 90 degrees (PI/2 radians).
                        zero_offset = 2048
                        
                        # 2048->0, 1024->90 3072->-90
                        
                        scale_factor = -(math.pi / 2.0) / 1024.0
                        radians = (pos - zero_offset) * scale_factor
                        if i == 5:
                            radians = math.pi/2.0 - radians
                      

                        # Use the correct joint names from the URDF file
                        urdf_joint_names = ["Rotation", "Rotation2", "Rotation3", "Rotation4", "Rotation5", "Rotation6"]
                        if i < len(urdf_joint_names):
                            raw_positions_dict[urdf_joint_names[i]] = radians
                    
                    latest_joint_data = raw_positions_dict
                    print("---")
                    time.sleep(0.1)

            except Exception as e:
                print(f"An error occurred: {e}")
                if bus.is_connected:
                    try:
                        bus.write('Torque_Enable', 0)
                    except Exception as e_write:
                        print(f"Could not disable torque: {e_write}")
                    bus.disconnect()
                print("Connection lost. Attempting to reconnect in 5 seconds...")
                time.sleep(5)

    except KeyboardInterrupt:
        print("\nStopping program.")
    finally:
        # Final cleanup when program exits (e.g., on Ctrl+C)
        if bus.is_connected:
            print("Disabling torque and disconnecting...")
            try:
                bus.write('Torque_Enable', 0)
            except Exception as e:
                print(f"Error disabling torque on exit: {e}")
            bus.disconnect()
            print("Disconnected.")

if __name__ == '__main__':
    main()

# endregion